## GANET V2
## Copyright (c) 2014 Ganet Development Team. All rights reserved.
##
## This file is subject to the terms and conditions defined in file
## 'LICENSE.md', which is part of this source code package.
## See 'AUTHORS' file for a list of contributors.

from .message import MessageEncoder, MessageDecoder

# Protocol identifier

_PROTOCOL_ID = 0x11E08200 # GA0200

# Network layer messages types and flags

_CONNECTION_REQUEST     = 0x01
_CONNECTION_ACCEPTED    = 0x02
_CONNECTION_REFUSED     = 0x03 
_DISCONNECT             = 0x04 
_HEARTBEAT              = 0x05 
_USER_DATA              = 0x06 

_SEQUENCED         = 0x10
_ORDERED           = 0x20
_RELIABLE          = 0x40


class RemoteAcknowledger(object):
    """Determines acknowledge information to be sent back to peer."""

    def __init__(self):        
        self.max_sqn = 0    # highest sequence number of remote
        self.ack_mask = 0   # acknowledge bit mask (32 bit)

    def __str__(self):
        return '%d + %s' % (self.max_sqn, bin(self.ack_mask))

    def acknowledge(self, remote_sqn):
        # the delta between highest sqn and given sqn determines if
        # we have to shift the bit mask and update the highest remote sqn
        delta = remote_sqn - self.max_sqn
        
        if delta > 0:
            # Shift bit mask to left by delta bits.
            
            # We need mask out values that exceed 32 bit
            # and just to keep it nice and neat we convert it back to int,
            # (might have become long int).
            self.ack_mask = int((self.ack_mask << delta) & 4294967295L) 

            # update acknowledge bit mask
            if delta <= 32 and self.max_sqn > 0:
                self.ack_mask |= 1 << (delta - 1)

            # update highest remote sequence number
            self.max_sqn += delta

            # message can be dispatched
            dispatch = True
        else:
            # update acknowledge bit mask
            if delta >= -32 and delta < 0:
                bit = 1 << (-delta - 1)
                # already dispatched?
                dispatch = not (self.ack_mask & bit) 
                self.ack_mask |= bit
            else:
                # message is to old to be dispatched
                dispatch = False

        return dispatch
    
class LocalAcknowledger(object):
    """Determinesn which of our sent messages got lost."""

    def __init__(self, deadline):
        self._messages = {}
        self._deadline = deadline
        self._max_sqn = 0           # highest sequence number of remote
        self._time = 0
        self._lost = []
        self._ack = []

    def add(self, sqn):
        assert sqn not in self._messages
        self._messages[sqn] = self._time + self._deadline

    def _acknowledge(self, sqn):
        if self._messages.pop(sqn, None):
            self._ack.append(sqn)

    def acknowledge(self, sqn, ack_mask):
        # update highest remote sequence number to remove lost messages due
        # to length of bit mask
        self._max_sqn = max(self._max_sqn, sqn)

        # remove current acknowledge message, if present
        self._acknowledge(sqn)

        # remove acknowledge messages due to bit mask, if present
        for i in range(31):
            if ack_mask & (1 << i):
                self._acknowledge(sqn - (i + 1))

    def _is_lost(self, msg_entry):
        return self._time >= msg_entry[1] or msg_entry[0] < self._max_sqn - 32

    def update(self, dt):
        self._time += dt

        lost = [x[0] for x in self._messages.iteritems() if self._is_lost(x)]
        for x in lost:
            self._messages.pop(x, None)
        self._lost.extend(lost)

    def get_acknowledge_messages(self):
        result = list(self._ack)
        self._ack[:] = []
        return result

    def get_lost_messages(self):
        result = list(self._lost)
        self._lost[:] = []
        return result


class OrderedMessageHandler(object):

    def __init__(self, peer):
        self._peer = peer
        self._next_sqn = 1
        self._messages = {}

    def _dispatch(self):
        msg = self._messages.pop(self._next_sqn, None)
        while msg:
            self._next_sqn += 1
            self._peer._process_by_type(msg[0], msg[1])
            msg = self._messages.pop(self._next_sqn, None)           

    def process_message(self, sqn, msg_type, msg):
        if sqn in self._messages or sqn < self._next_sqn:
            # discard message, already dispatched or stored
            return

        self._messages[sqn] = (msg_type, msg)  
        self._dispatch()
                
        
class Peer(object):

    def __init__(self, network, addr, heartbeat, timeout, dead = True):
        self._network = network
        self._addr = addr
        self._heartbeat = heartbeat
        self._timeout = timeout
        self._decoder = MessageDecoder()
        self._encoder = MessageEncoder()
        self._local_sqn = 1 # sqn '0' is reserved
        self._ordered_sqn = 1 # local next reliable ordered number
        self._max_seq = 0 # highest reliable sequenced number
        self._sequenced_sqn = 1 # local next reliable sequenced number
        self._time = 0
        self._headtbeat_deadline = 0
        self._timeout_deadline = self._timeout
        self._dead = dead
        
        # statistics
        self.messages_sent = 0
        self.messages_lost = 0
        
        # Currently it seems to be a good assumption that a message is
        # lost if not acknowledge after a duration of 3 heartbeats
        # This might be adapted after further testing.
        self._local_acknowledger = LocalAcknowledger(heartbeat * 3)
        self._remote_acknowledger = RemoteAcknowledger()
        self._ordered_message_handler = OrderedMessageHandler(self)

        # keer track of message to resent
        self._reliables = {}

    @property
    def dead(self):
        return self._dead

    @property
    def addr(self):
        return self._addr

    def _process_message(self, msg):
        # we got a message -> reset timeout deadline
        self._timeout_deadline = self._time + self._timeout
        
        # decode message header
        self._decoder.reset(msg)
        protocol_id = self._decoder.read_uint()
        
        if protocol_id != _PROTOCOL_ID:
            # we don't need to continued to decode the header any further
            # it seems not to be a valid message anyway
            raise ValueError('invalid protocol id %d' % protocol_id)

        sqn = self._decoder.read_uint()
        ack_sqn = self._decoder.read_uint()
        ack = self._decoder.read_uint()
        flag_n_type = self._decoder.read_byte()
        msg_type = flag_n_type & 0x0F
        flags = flag_n_type & 0xF0

        # process header
        self._local_acknowledger.acknowledge(ack_sqn, ack)
        if not self._remote_acknowledger.acknowledge(sqn):
            # we discard this message, already dispatched or to old
            self._network._log_debug('discarding message (#%d)' % sqn)
            return

        # process message type
        if flags & _ORDERED:
            assert flags & _RELIABLE
            ord_sqn = self._decoder.read_uint()
            user_data = self._decoder.get_remaining_message()
            #self._ordered_message_handler(ord_sqn, msg_type, user_data)
            self._ordered_message_handler.process_message(ord_sqn,
                                                          msg_type,
                                                          user_data)
        elif flags & _SEQUENCED:
            if flags & _RELIABLE:
                rsqn = self._decoder.read_uint()
                if rsqn > self._max_seq:
                    self._max_seq = rsqn
                    user_data = self._decoder.get_remaining_message()
                    self._process_by_type(msg_type, user_data)
                else:
                    self._network._log_debug('discarding message (#%d)' % sqn)
                    
            elif sqn == self._remote_acknowledger.max_sqn:
                user_data = self._decoder.get_remaining_message()
                self._process_by_type(msg_type, user_data)
        else:
            user_data = self._decoder.get_remaining_message()
            self._process_by_type(msg_type, user_data)

    def _process_by_type(self, msg_type, msg):
        if self._dead:
            if msg_type == _CONNECTION_REQUEST:
                ok = True
                for listener in self._network._listeners:
                    ok = ok and listener.on_connection_request(self)

                if ok:
                    self._dead = False
                    self._accept()
                    for listener in self._network._listeners:
                        listener.on_new_connection(self)
                else:
                    self._refuse()
            else:
                # peer is 'dead' ignore all other messages
                pass
        else:
            if msg_type == _HEARTBEAT:
                # nothing to do
                pass
            elif msg_type == _USER_DATA:
                # forward message to listeners
                for listener in self._network._listeners:
                    listener.on_message(self, msg)
            elif msg_type == _DISCONNECT:
                self.disconnect()
                # disconnect method will inform listeners
            elif msg_type == _CONNECTION_ACCEPTED:
                for listener in self._network._listeners:
                    listener.on_connection_accepted(self)
            elif msg_type == _CONNECTION_REFUSED:
                for listener in self._network._listeners:
                    listener.on_connection_refused(self)
                self._dead = True
            elif msg_type == _CONNECTION_REQUEST:
                # ignore multiple connection requests when connection is
                # already estapblished
                pass
            else:
                raise ValueError('unknown message type %d' % msg_type)
        

    def _send(self, msg):
        """Send message.
        The protocol header will be added to the specified message.
        The message must contain the message type and all required data.
        """

        # update local messages to be acknowledged        
        self._local_acknowledger.add(self._local_sqn)

        # create header
        self._encoder.reset()
        self._encoder.write_uint(_PROTOCOL_ID)
        self._encoder.write_uint(self._local_sqn)
        self._encoder.write_uint(self._remote_acknowledger.max_sqn)
        self._encoder.write_uint(self._remote_acknowledger.ack_mask)
        
        # queue message for sending
        self._network._send(self._addr, self._encoder.get_message() + msg)

        # update state
        self._heartbeat_deadline = self._time + self._heartbeat
        self._local_sqn += 1
        self.messages_sent += 1

    def _send_unreliable(self, msg_type, msg = ''):
        self._encoder.reset()
        self._encoder.write_byte(msg_type)
        self._send(self._encoder.get_message() + msg)

    def _send_unreliable_sequenced(self, msg_type, msg = ''):
        self._encoder.reset()
        self._encoder.write_byte(msg_type | _SEQUENCED)
        self._send(self._encoder.get_message() + msg)

    def _send_reliable(self, msg_type, msg = ''):
        self._encoder.reset()
        self._encoder.write_byte(msg_type | _RELIABLE)
        msg = self._encoder.get_message() + msg
        self._reliables[self._local_sqn] = msg
        self._send(msg)

    def _send_reliable_ordered(self, msg_type, msg = ''):
        self._encoder.reset()
        self._encoder.write_byte(msg_type | _RELIABLE | _ORDERED)
        self._encoder.write_uint(self._ordered_sqn)
        msg = self._encoder.get_message() + msg
        self._reliables[self._local_sqn] = msg
        self._send(msg)
        self._ordered_sqn += 1

    def _send_reliable_sequenced(self, msg_type, msg = ''):
        self._encoder.reset()
        self._encoder.write_byte(msg_type | _RELIABLE | _SEQUENCED)
        self._encoder.write_uint(self._sequenced_sqn)
        msg = self._encoder.get_message() + msg
        self._reliables[self._local_sqn] = msg
        self._send(msg)
        self._sequenced_sqn += 1
        
    def send_unreliable(self, msg):
        self._send_unreliable(_USER_DATA, msg)

    def send_unreliable_sequenced(self, msg):
        self._send_unreliable_sequenced(_USER_DATA, msg)

    def send_reliable(self, msg):
        self._send_reliable(_USER_DATA, msg)

    def send_reliable_ordered(self, msg):
        self._send_reliable_ordered(_USER_DATA, msg)

    def send_reliable_sequenced(self, msg):
        self._send_reliable_sequenced(_USER_DATA, msg)
    
    def disconnect(self):
        self._send_reliable(_DISCONNECT)
        self._dead = True
        
        for listener in self._network._listeners:
            listener.on_peer_disconnected(self)

    def _connect(self):
        self._send_reliable_ordered(_CONNECTION_REQUEST)

    def _refuse(self):
        self._send_reliable_ordered(_CONNECTION_REFUSED)

    def _accept(self):
        self._send_reliable_ordered(_CONNECTION_ACCEPTED)

    def _send_heartbeat(self):
        self._send_unreliable(_HEARTBEAT)

    def _handle_lost_message(self, lost_sqn):
        self._network._log_debug('lost message (sqn %d)' % lost_sqn)
        self.messages_lost += 1

        msg = self._reliables.get(lost_sqn, None)
        if msg:
            log_txt = 'resending reliable message (#%d) as (#%d)'
            log_txt = log_txt % (lost_sqn, self._local_sqn)
            self._network._log_debug(log_txt)
            self._reliables[self._local_sqn] = msg
            self._send(msg)
    
    def _update(self, dt):
        self._time += dt

        # local acknowledger will check for message deadlines
        self._local_acknowledger.update(dt)

        # handle lost messages
        for lost_sqn in self._local_acknowledger.get_lost_messages():
            self._handle_lost_message(lost_sqn)

        # Remove acknowledge reliable messages
        for ack_sqn in self._local_acknowledger.get_acknowledge_messages():
            self._reliables.pop(ack_sqn, None)

        if self._dead:
            return
        
        # check for heartbeat
        if self._time >= self._heartbeat_deadline:
            self._send_heartbeat()

        # check for timeout
        if self._time >= self._timeout_deadline:
            self.disconnect()
            for listener in self._network._listeners:
                listener.on_peer_lost(self)
